package tree.structure.trie;

/**
 * 前缀树的数组实现形式
 */
public class Trie2 {

    private int[] passArr;
    private int[] endArr;
    private int maxStrLen;
    private int cnt = 1;
    private int[][] t;

    public Trie2(int maxStrLen) {
        passArr = new int[maxStrLen + 1];
        endArr = new int[maxStrLen + 1];
        this.maxStrLen = maxStrLen;
        t = new int[maxStrLen + 1][26];
    }

    public void insert(String s) {
        passArr[0]++;
        char[] arr = s.toCharArray();
        int tempCnt = 1;
        for (char ch : arr) {
            int idx = ch - 'a';
            if (t[tempCnt][idx] == 0) {
                t[tempCnt][idx] = ++cnt;
            }
            tempCnt = t[tempCnt][idx];
            passArr[tempCnt]++;
        }
        endArr[tempCnt]++;
    }

    public int countWord(String s) {
        int tempCnt = 1;
        char[] chars = s.toCharArray();
        for (char ch : chars) {
            int idx = ch - 'a';
            if (t[tempCnt][idx] == 0) {
                return 0;
            }
            tempCnt = t[tempCnt][idx];
        }

        return endArr[tempCnt];
    }

    public int countPrefixWord(String s) {
        int tempCnt = 1;
        char[] chars = s.toCharArray();
        for (char ch : chars) {
            int idx = ch - 'a';
            if (t[tempCnt][idx] == 0) {
                return 0;
            }
            tempCnt = t[tempCnt][idx];
        }

        return passArr[tempCnt];
    }

    public void erase(String s) {
        if (countWord(s) < 1) {
            return;
        }
        int tempCnt = 1;
        passArr[0]--;

        char[] chars = s.toCharArray();
        for (char ch : chars) {
            int idx = ch - 'a';
            passArr[t[tempCnt][idx]]--;
            if (passArr[t[tempCnt][idx]] == 0) {
                t[tempCnt][idx] = 0;
                return;
            }
            tempCnt = t[tempCnt][idx];
        }
        
        endArr[tempCnt]--;
    }

    public static void main(String[] args) {
        Trie2 trie = new Trie2(10);
        trie.insert("a");
        trie.insert("abc");
        trie.insert("ad");
        trie.insert("ab");
        trie.insert("ab");

//        trie.erase("ab");

        System.out.println("trie.countWord(\"a\") = " + trie.countWord("a"));
        System.out.println("trie.countWord(\"aa\") = " + trie.countWord("aa"));
        System.out.println("trie.countWord(\"ab\") = " + trie.countWord("ab"));

        System.out.println("trie.countPrefixWord(\"a\") = " + trie.countPrefixWord("a"));

        System.out.println("trie = " + trie);
    }
}
